#!/usr/bin/env python3
import os
import sys
import re

"""
This script removes:
1. Citation links in format [[number]](url)
2. Markdown blockquotes starting with >
3. Code blocks marked with ```note
Can process either a single markdown file or all markdown files in a directory.
"""
def replace_in_file(file_path):
    if not file_path.endswith('.md'):
        print(f"Error: {file_path} is not a markdown file")
        return

    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()

        # Remove citation links [[number]](url)
        content = re.sub(r'\[\[\d+\]\]\(.*?\)', '', content)
        content = re.sub(r'\[\[\d+\\?\]\]\(https://poe\.com/citation\?[^)]+\)', '', content)


        # Remove blockquotes (lines starting with >)
        content = re.sub(r'^>.*$', '', content, flags=re.MULTILINE)

        # Remove code blocks marked with ```note
        content = re.sub(r'```note\n[\s\S]*?```', '', content)

        # Remove empty lines that might be created after removal
        content = re.sub(r'\n\s*\n', '\n\n', content)

        file_output_path = file_path.replace('.md', '-final.md')

        with open(file_output_path, 'w', encoding='utf-8') as file:
            file.write(content)
        print(f"Processed: {file_path} -> {file_output_path}")

    except FileNotFoundError:
        print(f"Error: File {file_path} not found")
    except Exception as e:
        print(f"Error processing {file_path}: {str(e)}")

def process_directory(directory):
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.md'):
                file_path = os.path.join(root, file)
                replace_in_file(file_path)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python script.py <file_path_or_directory>")
        sys.exit(1)

    path = sys.argv[1]

    # Check if path exists
    if not os.path.exists(path):
        print(f"Error: {path} does not exist")
        sys.exit(1)

    # Process based on whether path is file or directory
    if os.path.isfile(path):
        replace_in_file(path)
    elif os.path.isdir(path):
        process_directory(path)
    else:
        print(f"Error: {path} is neither a valid file nor directory")
        sys.exit(1)

    print("Processing complete.")

# Usage:
# For single file: python3 remove_ai.py /path/to/your/file.md
# For directory: python3 remove_ai.py /path/to/your/markdown/files
